from vector_database import database, get_files
import os

def init(drop_collection: bool = False):
    client = database.get_client()
    if client.has_collection(collection_name=os.getenv("MILVUS_COLLECTION")) and drop_collection == True:
        client.drop_collection(collection_name=os.getenv("MILVUS_COLLECTION"))
    client.close()

    file_list = get_files.get_file_list()

    count = {}

    for file in file_list:
        print(f"\nAdding {file} to database:\n")
        res = database.insert_db(database.get_records(filename=file))
        count[get_files.create_collection_name(file)] = res["insert_count"]

    print("\n")
    for i, r in enumerate(count):
        print(f"{i}: Insert {count[r]} records from the file '{r}'")
